package router

import (
	"github.com/gin-gonic/gin"
	"io"
	"log"
	"net/http"
	"github.com/go-resty/resty/v2"
)

type SsrfHandler struct {
}

func NewGetSsrfHandler() *SsrfHandler {
	return &SsrfHandler{}
}

func (handler *SsrfHandler) RegisterRoutes(r *gin.RouterGroup) {
	r.GET("/ssrf", handler.Ssrf)
	r.GET("/ssrf_sec", handler.SsrfSec)
	r.GET("/ssrf_resty", handler.SsrfWithResty)
	r.GET("/ssrf_resty_sec", handler.SsrfRestySec)
}

func (handler *SsrfHandler) Ssrf(ctx *gin.Context) {
	url := ctx.Query("url")
	resp, err := http.Get(url)
	if err != nil {
		ctx.String(http.StatusBadRequest, "Error: %v", err)
		return
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)

	ctx.String(http.StatusOK, string(body))
}

func (handler *SsrfHandler) SsrfSec(ctx *gin.Context) {
	host := ctx.Query("host")

	// 检查host是否在白名单中
	if !isWhiteListHost(host) {
		ctx.String(http.StatusBadRequest, "Invaild host")
		return
	}

	client := &http.Client{
		// 禁止重定向
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	resp, err := client.Get("http://" + host + ":8080/test/redirect")	//协议硬编码
	if err != nil {
		ctx.String(http.StatusBadRequest, "Error: %v", err)
		return
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		ctx.String(http.StatusBadRequest, "Read body err: %v", err)
		return
	}

	respBody := string(body)
	log.Printf("respBody: %v\n", respBody)	//尽量避免将远程响应返回给接口响应，这里只输出日志

	ctx.Status(http.StatusOK)
}

func (handler *SsrfHandler) SsrfWithResty(ctx *gin.Context) {
	url := ctx.Query("url")
	client := resty.New()
	// 禁止自动处理重定向
	client.SetRedirectPolicy(resty.NoRedirectPolicy())
	resp, err := client.R().SetHeader("User-Agent", "resty").Get(url)

	if err != nil {
		ctx.String(http.StatusBadRequest, "Error: %v", err)
		return
	}

	ctx.String(http.StatusOK, resp.String())
}

func (handler *SsrfHandler) SsrfRestySec(ctx *gin.Context) {
	host := ctx.Query("host")

	// 检查host是否在白名单中
	if !isWhiteListHost(host) {
		ctx.String(http.StatusBadRequest, "Invaild host")
		return
	}

	client := resty.New()
	// 禁止自动处理重定向
	client.SetRedirectPolicy(resty.NoRedirectPolicy())
	resp, err := client.R().
		SetHeader("User-Agent", "resty").
		Get("http://" + host + ":8080/test/redirect")	//协议硬编码

	if err != nil {
		ctx.String(http.StatusBadRequest, "Error: %v", err)
		return
	}

	log.Printf("resp: %v\n", resp.String())
	ctx.Status(http.StatusOK)
}

func isWhiteListHost(host string) bool {
	//设置host白名单
	whiteListHost := []string{"172.30.14.55", "example.com"}
	for _, whiteHost := range whiteListHost {
		if host == whiteHost {
			return true
		}
	}
	return false
}